from langchain_openai import ChatOpenAI
from langchain_core.embeddings import Embeddings
from ai_configs import defalut_ai_config
import requests
import json


def create_llm(code='doubao', temperature=0.5):
    _ai_config = defalut_ai_config[code]

    if _ai_config is None:
        raise Exception("未传入配置信息")

    return ChatOpenAI(
        base_url=_ai_config.get("url").replace("/chat/completions", ""),
        api_key=_ai_config.get("key"),
        model=_ai_config.get("model"),
        temperature=temperature,
    )


def create_embeddings(platform_code="doubao-embedding"):
    _ai_config = defalut_ai_config[platform_code]

    if _ai_config is None:
        raise Exception("未传递配置参数信息")

    return CustomEmbeddings(
        base_url=_ai_config.get("url").replace('/embeddings', ""),
        api_key=_ai_config.get("key"),
        model=_ai_config.get("model"),
    )


class CustomEmbeddings(Embeddings):
    def __init__(self, base_url, api_key, model):
        self.base_url = base_url
        self.api_key = api_key
        self.model = model

    def embed_documents(self, texts):
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }
        payload = {"model": self.model, "input": texts, "encoding_format": "float"}

        reponse = requests.post(
            f"{self.base_url}/embeddings", headers=headers, data=json.dumps(payload)
        )

        # 检查请求是否成功，如果失败则抛出错误
        reponse.raise_for_status()

        # 解析响应JSON数据
        json_data = reponse.json()

        # 从响应数据中提取嵌入向量并返回
        return [item["embedding"] for item in json_data["data"]]

    def embed_query(self, text):
        return self.embed_documents([text])[0]
